#include "MallocChest.h"
#include <assert.h>
#include <new>
#include <iostream>
#include "../Utils/LogUtil.h"
#include "../Templates/Alignmen.h"
#include "../Templates/TypeCompatibleBytes.h"
#include "../Utils/PlatTlsUtil.h"

struct PtrChunk
{
	const uint32 mMask;
	const uint32 mAllocationNum;
	uint32 mNumUsed;
	PtrChunk* mNext;
	PtrData* mAllocations[MallocChest::MAX_NUM_ALLOCS_IN_POOL];
	uint8* mPool;

	PtrChunk(const uint32 mask, const uint32 numAllocs) :
		mMask(mask),
		mAllocationNum(numAllocs),
		mNumUsed(0),
		mNext(nullptr)
	{
		mPool = (uint8*)malloc((uint64)(mMask + MallocChest::PTRDATA_OVERHEAD) * mAllocationNum);
		uint32 offset = 0;
		for (uint32 index = 0; index < mAllocationNum; ++index)
		{
			uint8* newPtr = mPool + offset;
			mAllocations[index] = new(newPtr) PtrData(mMask, newPtr + MallocChest::PTRDATA_OVERHEAD);
			offset += mMask;
			offset += MallocChest::PTRDATA_OVERHEAD;
		}
	}
	~PtrChunk()
	{
		free(mPool);
	}

	uint8* AllocateFromChunk(uint32 allocationSize)
	{
		PtrData* ptrInfo = nullptr;
		for (uint32 index = 0; index < mAllocationNum; ++index)
		{
			PtrData* ptrIt = mAllocations[index];
			if (ptrIt->mSize == 0)
			{
				ptrInfo = ptrIt;
				break;
			}
		}

		assert(ptrInfo != nullptr);

		mNumUsed++;
		ptrInfo->mSize = allocationSize;

		memset((void*)ptrInfo->mPtr, MallocChest::MEM_TAG, allocationSize);

		// LOGI(L"Malloc Size=%u, PooledPtr=0x%016llx", allocationSize, (uint64_t)ptrInfo->mPtr);

		return ptrInfo->mPtr;
	}

	void TryFreeFromChunk(uint8* ptr)
	{
		bool isRemoved = false;
		for (uint32 index = 0; index < mAllocationNum; ++index)
		{
			PtrData* ptrIt = mAllocations[index];
			if (ptrIt->mPtr == ptr)
			{
				memset((void*)ptrIt->mPtr, MallocChest::MEM_WIPETAG, ptrIt->mSize);

				// LOGI(L"Free Size=%u, PooledPtr=0x%016llx", ptrIt->mSize, (uint64_t)ptr);

				ptrIt->mSize = 0;
				isRemoved = true;
				mNumUsed--;

				break;
			}
		}

		assert(isRemoved);
	}

	bool IsPtrInChunk(uint8* ptr)
	{
		return (ptr >= mPool && ptr < mPool + (uint64)(mMask + MallocChest::PTRDATA_OVERHEAD) * mAllocationNum);
	}
};

MallocChest::MallocChest()
{

	mAllChunkDesc = {
		{64, 224},
		{96, 144},
		{128, 256},
		{192, 512},
		{256, 384},
		{384, 208},
		{512,  48},
		{768,  32},
		{1024,  32},
		{2048,  32},
		{4096,  32},
		{8192,  32},
		{16384,  16},
		{32768,  16}
	};

	for (auto& it : mAllChunkDesc)
	{
		mMemoryChunks.insert({ it.first, new PtrChunk(it.first, it.second) });
	}

	mLargePool = (uint8*)malloc(LARGE_MEMORYPOOL_SIZE + PTRNODE_OVERHEAD * 2);
	assert(mLargePool != nullptr);

	mLargeTail = (PtrNode*)Align(mLargePool + LARGE_MEMORYPOOL_SIZE + PTRNODE_OVERHEAD, PTRNODE_ALIGNMENT);
	mLargeTail->mSize = 0;
	mLargeTail->mNext = nullptr;
	mLargeTail->mPtr = nullptr;

	mLargeHead = (PtrNode*)Align(mLargePool, PTRNODE_ALIGNMENT);
	mLargeHead->mSize = 0;
	mLargeHead->mNext = mLargeTail;
	mLargeHead->mPtr = (uint8*)(mLargeHead) + PTRNODE_OVERHEAD;


}

MallocChest::~MallocChest()
{
	for (auto& it : mMemoryChunks)
	{
		delete it.second;
	}

	free(mLargePool);
}

void* MallocChest::Malloc(uint32 count, uint32 alignment /*= DEFAULT_ALIGNMENT*/)
{
	// if (count <= mAllChunkDesc.end()->first)
	if(count <= 32768)
	{
		uint32 mask = 0;
		for (auto& it : mAllChunkDesc)
		{
			if (count < it.first)
			{
				mask = it.first;
				break;
			}
		}
		if (mask > 0)
		{
			PtrChunk* chunk = mMemoryChunks.find(mask)->second;
			assert(chunk != nullptr);
			while (chunk != nullptr)
			{
				if (chunk->mNumUsed < chunk->mAllocationNum)
					break;
				chunk = chunk->mNext;
			}
			if (chunk != nullptr)
			{
				return chunk->AllocateFromChunk(count);
			}
			else
			{
				chunk = mMemoryChunks.find(mask)->second;
				while (chunk->mNext)
					chunk = chunk->mNext;
				chunk->mNext = new PtrChunk(mask, mAllChunkDesc.find(mask)->second);
				return chunk->mNext->AllocateFromChunk(count);
			}
		}
		else
		{
			assert(0);
		}
	}
	else
	{
		// use big memory
		PtrNode* p = mLargeHead;
		PtrNode* prev = nullptr;
		// mLargeHead.size always 0
		uint32 minSpace = static_cast<uint32>((uint8*)mLargeTail - mLargeHead->mPtr);
		// uint32 minSpace = static_cast<uint32>((uint8*)mLargeTail - Align(mLargeHead->mPtr + mLargeHead->mSize, PTRNODE_ALIGNMENT));
		while (p->mNext)
		{
			uint32 currSpace = static_cast<uint32>((uint8*)p->mNext - Align(p->mPtr + p->mSize, PTRNODE_ALIGNMENT));
			if (currSpace <= minSpace && currSpace >= count + PTRNODE_OVERHEAD)
			{
				minSpace = currSpace;
				prev = p;
			}
			p = p->mNext;
		}

		if (prev == nullptr)
		{
			// No space for malloc
			assert(0);
		}
		
		PtrNode* result = (PtrNode*)Align(prev->mPtr + prev->mSize, PTRNODE_ALIGNMENT);
		result->mSize = count;
		result->mNext = prev->mNext;
		prev->mNext = result;
		result->mPtr = (uint8*)result + PTRNODE_OVERHEAD;
		memset(result->mPtr, MEM_TAG, count);
		return result->mPtr;
	}
	return nullptr;
}

void* MallocChest::Realloc(void* original, uint32 count, uint32 alignment /*= DEFAULT_ALIGNMENT*/)
{
	void* result = nullptr;
	if (IsPtrInLarge((uint8*)original))
	{
		PtrNode* ptrNode = (PtrNode*)((uint8*)original - PTRNODE_OVERHEAD);

		result = Malloc(count, alignment);
		memcpy(result, original, ptrNode->mSize);
		Free(original);
	}
	else
	{
		PtrData* ptrInfo = (PtrData*)((uint8*)original - PTRDATA_OVERHEAD);
		if (ptrInfo->mMask >= count)
		{
			return original;
		}

		result = Malloc(count, alignment);
		memcpy(result, original, ptrInfo->mSize);
		Free(original);
	}
	return result;
}

void MallocChest::Free(void* original)
{
	if (IsPtrInLarge((uint8*)original))
	{
		PtrNode* ptrNode = (PtrNode*)((uint8*)original - PTRNODE_OVERHEAD);
		PtrNode* p = mLargeHead;
		PtrNode* pre = nullptr;
		while (p->mNext)
		{
			if (p->mNext == ptrNode)
			{
				pre = p;
				break;
			}
			p = p->mNext;
		}
		assert(pre != nullptr);
		pre->mNext = ptrNode->mNext;
	}
	else
	{
		PtrData* ptrInfo = (PtrData*)((uint8*)original - PTRDATA_OVERHEAD);
		PtrChunk* chunk = mMemoryChunks.find(ptrInfo->mMask)->second;
		assert(chunk != nullptr);
		while (chunk)
		{
			if (chunk->IsPtrInChunk((uint8*)original))
				break;
			chunk = chunk->mNext;
		}
		assert(chunk != nullptr);
		chunk->TryFreeFromChunk((uint8*)original);
	}
}

bool MallocChest::IsPtrInLarge(uint8* ptr)
{
	return (ptr >= (uint8*)mLargeHead->mPtr && ptr < (uint8*)mLargeTail);
}

